#include "light_path.h"
#include "scene.h"
//#define LOGGING_ON
#ifdef LOGGING_ON
#include <fstream>
using std::endl;
namespace gemren
{
	std::ofstream logfile("caustics_log.txt");
}
#endif
namespace gemren

{
	void light_paths::compute_lighting(const scene* sc, random* ran, intersection* i) const 
	{
		size_t mytri = i->id0;

		if(!totalcnts[mytri]) return; //no hits = no caustic lighting
		
		//choose my transmitting triangle
		int hitn = ran->get()%totalcnts[mytri];
		size_t choice = tripairs[mytri];
		while(hitn >= int(hitcounts[choice]))
		{
			hitn -= int(hitcounts[choice]);
			++choice;
		}
		
		const triangle& tri = sc->get_triangle(transmitters[choice]);
		
		coord2d samcoord = ran->uniform_triangle();
		const coord& v0 = sc->get_vertex(tri.v0);
		intersection li; //copy
		li.abs_pos = i->abs_pos;
		li.set_wavelength( i->get_wavelength());
		coord outcoming =  (v0 + (sc->get_vertex(tri.v1) - v0)*samcoord.x + (sc->get_vertex(tri.v2) - v0)*samcoord.y) -i->abs_pos;
		float dist = embree::length(outcoming);
		
		outcoming /= dist;
		li.outcoming = outcoming; 

		int depth = 7;
		
		float pdf = (float(hitcounts[choice]))/float(totalcnts[mytri])  //pdf of choosing transmitter
			          *dist*dist / (embree::clamp(-embree::dot(tri.normal, outcoming))*tri.area(sc)); //pdf of choosing direction
		float factor;
		if(i->lighting_mode & LM_CAUSTIC_MIS)
			factor = 1/(pdf+ i->triptr->insidemat->pdf(i->incoming, outcoming, i->normal)); //*pdf/pdf
		else /* LM_CAUSTIC_LIGHTSOURCE */ 
			factor = 1/pdf;
		while(true)//--depth)
		{
			if(!sc->first_intersection(li, li.abs_pos, li.outcoming, EPSILON)) return;
			
			if(li.triptr->insidemat->is_light_emitting())
			{
				if(embree::dot(li.triptr->normal, li.outcoming)>0) return; //hit light from the other side - can happen
				light_emitting_material* lmat = (light_emitting_material*)(li.triptr->insidemat);
				li.add_light(lmat->get_intensity()*i->triptr->insidemat->throughput(outcoming, -i->incoming, i->normal)* factor);
				i->add_light( li );
				
				return;
			}
			if (!li.triptr->insidemat->is_light_transmitting()) return;
			else li.triptr->insidemat->bounce(ran, sc, &li);
			
			float rr = embree::clamp(li.reflectance, 0.3f, 0.8f);
			if(li.reflectance == 0 || ran->getf()>rr) break;
			else li.dim_light(1.0f/rr);
		}
		
	}
	
	
	void light_paths::search_for_paths(::size_t tid, void* ptr, ::size_t elt)
	{
		chunkdata* d = (chunkdata*)ptr;
		size_t cnt = d->ray_count;
		scene* sc = d->sc;
		map_t paths;
		random r = d->rand[elt];
		

		for(size_t i = 0; i < cnt; ++i)
		{
			int weight = 1;
			intersection isec;
			isec.set_wavelength(r.getf());
			std::pair<size_t, size_t> pr;
			size_t lid;
			if((lid = sc->get_light_triangle(r) == size_t(-1))) break;
			if(d->candidates->empty()) break;
			transmitter(pr) = lid;

			const triangle& tri = sc->get_triangle(lid);
		
			
			isec.abs_pos = sc->triangle_area_sample(r, lid);

			size_t dest = d->candidates->at(r.get()%(d->candidates->size()));
			isec.outcoming = sc->triangle_area_sample(r, dest) - isec.abs_pos;
			if(embree::dot(isec.outcoming, tri.normal)<0.0f) continue;
			isec.outcoming = embree::normalize(isec.outcoming);
			int depth = 0;
			while( sc->first_intersection(isec, isec.abs_pos, isec.outcoming, EPSILON))
			{
				material* mat = isec.triptr->insidemat;
				if(mat->is_directly_lighted())
				{
					if(depth == 0) break; //= normal direct lighting
					receiver(pr) = isec.id0;
					
					std::pair<map_t::iterator, bool> res =  paths.insert( std::pair<rtpair, size_t>(pr,weight));
					if(!res.second) //was there already
					{
						hitcount(res.first)+=weight;
					}
					break;
				}
				else if(mat->is_light_transmitting())
				{
					transmitter(pr) = isec.id0;
					mat->interact(&r, sc, &isec);
					
					//play russian roulette
					float rr = embree::clamp(isec.reflectance, 0.01f, 0.99f);
					if(r.getf()>rr) break;
					
					++depth;
				}
				else break;
			}
		}
		d->rand[elt] = r;
		d->sets[elt].swap(paths);
	}
	
	
	void light_paths::find_paths(scene* sc)
	{
		ensure_scheduler_init();
		
		size_t threads = embree::scheduler->getNumThreads();
		
		d.reset(threads);
		d.ray_count = iteration_size/threads;
		d.sc = sc;
		d.candidates = &transmitting_candidates;
		get_transmitting_candidates(sc);
		int iter = 0;
		do
		{
			
#ifdef LOGGING_ON
			logfile<<"iteration: "<<iter+1<<endl;
#endif
			for(size_t i = 0; i< threads; ++i) d.sets[i].clear();
			embree::scheduler->addTask(search_for_paths, &d, threads);
			embree::scheduler->go();
		}
		while(merge_paths(d.sets, threads)>0 && ++iter < max_iterations);
		inform_triangles(sc);
	}

	
	void light_paths::get_transmitting_candidates(const scene* sc)
	{
		size_t ind = 0;
		for(const triangle* tr = sc->tbegin(); tr != sc->tend(); ++tr)
		{
			if(tr->insidemat->is_light_transmitting()) 
				transmitting_candidates.push_back(ind);
			++ind;
		}
	}

	size_t light_paths::merge_paths(const map_t* sets, const size_t count)
	{ //add hit counts to single set <pairs>
		size_t added = 0;
		for(size_t i = 0; i< count; ++i)
		{
			const map_t& np(sets[i]);
			for(map_t::const_iterator i = np.begin(); i!= np.end(); ++i)
			{
				std::pair<map_t::iterator, bool> res = pairs.insert(*i);
				if(!res.second) //is there already
				{
					hitcount(res.first) += hitcount(i);
				}
				else ++added;
			}
		}
			
#ifdef LOGGING_ON
			logfile<<"added: "<<added<<" size: "<<pairs.size()<<endl;
#endif

		return added;
	}
	void light_paths::finish_caustic_mis_coef(intersection* i, const scene* sc) const
	{   
		//this method is only applicable with caustic MIS
		//is called when brdf sampled ray from diffuse mat hit transmitting mat
		//now is the time to compute mis weight for brdf sampled light that the ray may hit later on
		size_t receiver = i->prevtri, 
			transmitter = i->id0;
		size_t begin = tripairs[receiver];
		size_t total = totalcnts[receiver];
		if(!total)
		{
			i->mis_brdf_importance = 1.0f; //no registered transmitters
			return;
		}
		size_t* cur_trans = transmitters+begin;
		size_t* cur_hits = hitcounts+begin;
		size_t rem = total;
		while(rem > 0)
		{
			if(*cur_trans == transmitter)
			{
				float coslo = embree::clamp(-embree::dot(i->triptr->normal, i->incoming));
				float ls_importance = i->t * i->t * *cur_hits / (total * i->triptr->area(sc)*coslo ); //cur_hits/total = probability of selecting this transmitter
				i->mis_brdf_importance = i->mis_brdf_importance / (i->mis_brdf_importance + ls_importance);
				return;
			}
			rem -= *cur_hits;
			++cur_trans;
			++cur_hits;
		}
		i->mis_brdf_importance = 1.0f; //triangle that was hit is not on the transmitters list - even if its caustic, it still isnt lightsource sampled
	}

	void light_paths::inform_triangles(scene* sc)
	{
		delete[] tripairs;
		delete[] totalcnts;
		delete[] transmitters;
		delete[] hitcounts;
		size_t tcnt = sc->tri_count();
		tripairs = new size_t[tcnt];
		totalcnts = new size_t[tcnt];
		transmitters = new size_t[pairs.size()];
		hitcounts = new size_t[pairs.size()];

		for(size_t i = 0; i< tcnt; ++i)
		{
			totalcnts[i] = 0;
			tripairs[i] = 0;
		}
		for( map_t::iterator it = pairs.begin(); it!= pairs.end(); ++it)
		{
			if(hitcount(it)<threshold) continue;
			++tripairs[receiver(it)];
		}
		for(size_t i = 1; i< tcnt; ++i) 
		{
			tripairs[i] += tripairs[i-1]; //totalcnt[i] is now pointer at end
		}
		size_t totalc= 0, totalp=0;
		for( map_t::iterator it = pairs.begin(); it!= pairs.end(); ++it)
		{
			if(hitcount(it)<threshold) continue;
			size_t trpos = --tripairs[ receiver(it) ];
		
			transmitters[ trpos ] = transmitter(it);
			
			size_t weightedcount = size_t(hitcount(it) * -1.0f*embree::dot(sc->get_triangle( transmitter(it)).normal, sc->get_triangle(receiver(it)).normal)); 
			
			hitcounts[ trpos] = hitcount(it);
			totalcnts[ receiver(it) ] += hitcount(it);
			++totalp;
			totalc += hitcount(it);
		}
		
		pairs.clear();
#ifdef LOGGING_ON
		for(int i = 0; i< tcnt; ++i)
			logfile<<"t "<<i<<" has "<<tripairs[i]<<" paths,\t"<<totalcnts[i]<<" hits."<<endl;
		logfile<<"total "<<totalp<<" paths,\t"<<totalc<<" hits."<<endl;

#endif
	}

	err_t light_paths::set_my_int(std::string& name, const int value)
	{
		if(name == "rays") iteration_size = value;
		else if(name == "iterations") max_iterations = value;
		else if(name == "threshold") threshold = value;
		else	return ERR_BAD_PROP;
		return ERR_OK;
	}
	void light_paths::announce_properties(const std::string& prefix, property_announcer announcer)
	{
		announcer((prefix+".rays").c_str(), ANNOUNCED_INT, val_to_str(iteration_size).c_str());

		//announcer((prefix+".iterations").c_str(), ANNOUNCED_INT, val_to_str(max_iterations).c_str());
		//announcer((prefix+".threshold").c_str(), ANNOUNCED_INT, val_to_str(threshold).c_str());
	}
}